using System;
using System.Collections.Generic;
using System.Linq;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics;
using Atomic.Structures;

namespace Atomic.Thermodynamics.PartitionFunctions
{
	[Serializable]
	public class PartitionFunction : IHelmholtzGrandPotential
	{
		private Atom atom;
		private double[] counts;
		private double count_min, count_max;
		private RollingCache<Tuple<double, double>, InnerSolver> solvers;
		private RollingCache<double, Tuple<double, double>> volumeRanges;

		public PartitionFunction(IEnumerable<IPartitionFunctionEntry> entries)
		{
			if (entries == null)
			{
				throw new ArgumentNullException();
			}

			Entries = new ReadOnlyList<IPartitionFunctionEntry>(entries);

			if (Entries.Count == 0)
			{
				throw new ArgumentException();
			}

			Atom[] atoms = entries.SelectMany(e => e.Composition.Atoms).Distinct().Where(a => a != Atom.Vacancy).ToArray();

			if (atoms.Length == 0)
			{
				throw new ArgumentException();
			}

			if (atoms.Length != 1)
			{
				throw new NotSupportedException("Currently only 1-atom systems are supported.");
			}

			// The non-vacancy atom and the corresponding per-unit cell counts.
			atom = atoms[0];
			counts = entries.Select(p => p.Composition.Count(atom)).ToArray();
			count_min = counts.Min();
			count_max = counts.Max();

			// Rolling cache of solvers different temperatures and volumes.
			solvers = new RollingCache<Tuple<double, double>, InnerSolver>(10000);

			// And for volume search range for different temperatures.
			volumeRanges = new RollingCache<double, Tuple<double, double>>(10000);

			GibbsPotential = new InnerGibbsPotential(this);
		}

		public double FreeEnergy(double temperature, double volume, Composition composition)
		{
			double x = Count(composition);
			return CreateSolver(temperature, volume).HelmholtzPotential(x);
		}

		public double Pressure(double temperature, double volume, Composition composition)
		{
			double x = Count(composition);
			return CreateSolver(temperature, volume).Pressure(x);
		}

		public ChemicalPotential ChemicalPotential(double temperature, double volume, Composition composition)
		{
			double x = Count(composition);
			return new ChemicalPotential(atom, CreateSolver(temperature, volume).ChemicalPotential(x));
		}

		public double PressureEquilibriumVolume(double temperature, double pressure, Composition composition)
		{
			if (temperature < 0.0)
			{
				throw new ArgumentOutOfRangeException();
			}

			double x = Count(composition);

			if (x < count_min || x > count_max)
			{
				// This can't be reached.
				return double.NaN;
			}

			double v0, v1;

			Tuple<double, double> volumeRange;
			if (volumeRanges.TryGetValue(temperature, out volumeRange))
			{
				v0 = volumeRange.Item1;
				v1 = volumeRange.Item2;
			}
			else
			{
				v0 = double.PositiveInfinity;
				v1 = double.NegativeInfinity;

				foreach (IPartitionFunctionEntry e in Entries)
				{
					double v = e.Potential.EquilibriumVolume(temperature, pressure);
					v0 = Math.Min(v0, v);
					v1 = Math.Max(v1, v);
				}

				volumeRanges.Add(temperature, Tuple.Create<double, double>(v0, v1));
			}

			//return RootFinder.Bisection(v => CreateSolver(temperature, v).Pressure(x) - pressure, vmin, vmax);
			return RootFinder.Bisection(v => CreateSolver(temperature, v).Pressure(x) - pressure, v0, v1);
		}
		
		private double Count(Composition composition)
		{
			if (composition == null)
			{
				throw new ArgumentNullException();
			}

			foreach (Atom a in composition.Atoms)
			{
				if (a == Atom.Vacancy || a == Atom.Undefined)
				{
					// Ignore vacancies.
					continue;
				}

				if (a != atom)
				{
					throw new ArgumentException();
				}
			}

			return composition.Count(atom);
		}

		private InnerSolver CreateSolver(double temperature, double volume)
		{
			if (temperature < 0.0 || volume <= 0.0)
			{
				throw new ArgumentOutOfRangeException();
			}

			InnerSolver solver;
			if (!solvers.TryGetValue(Tuple.Create<double, double>(temperature, volume), out solver))
			{
				solver = new InnerSolver(this, temperature, volume);
				solvers.Add(Tuple.Create<double, double>(temperature, volume), solver);
			}

			return solver;
		}

		public ReadOnlyList<IPartitionFunctionEntry> Entries
		{
			get;
			private set;
		}

		public IGibbsGrandPotential GibbsPotential
		{
			get;
			private set;
		}

		[Serializable]
		private class InnerGibbsPotential : IGibbsGrandPotential
		{
			private PartitionFunction parent;

			public InnerGibbsPotential(PartitionFunction parent)
			{
				this.parent = parent;
			}

			public double FreeEnergy(double temperature, double pressure, Composition composition)
			{
				double volume = Volume(temperature, pressure, composition);

				// Legendre transform. The first term is the Helmholtz free energy.
				return parent.FreeEnergy(temperature, volume, composition) + pressure * volume;
			}

			public double Volume(double temperature, double pressure, Composition composition)
			{
				return parent.PressureEquilibriumVolume(temperature, pressure, composition);
			}

			public ChemicalPotential ChemicalPotential(double temperature, double pressure, Composition composition)
			{
				double volume = Volume(temperature, pressure, composition);
				return parent.ChemicalPotential(temperature, volume, composition);
			}

			public double VolumetricThermalExpansionCoefficient(double temperature, double pressure, Composition composition)
			{
				throw new NotImplementedException();
			}

			public IHelmholtzGrandPotential HelmholtzPotential
			{
				get
				{
					return parent;
				}
			}
		}

		[Serializable]
		private class InnerSolver
		{
			private PartitionFunction parent;
			private int n;
			private double[] energies, pressures;
			private double beta, tol, mu0, mu1;

			public InnerSolver(PartitionFunction parent, double t, double v)
			{
				this.parent = parent;

				n = parent.Entries.Count;

				// Precompute free energies for the given fixed temperature and volume.
				energies = new double[n];
				pressures = new double[n];
				for (int i = 0; i < n; i++)
				{
					energies[i] = parent.Entries[i].Potential.FreeEnergy(t, v);
					pressures[i] = parent.Entries[i].Potential.Pressure(t, v);
				}

				double kb = 8.617332478e-5; // eV/K
				beta = 1.0 / (kb * t);

				// Determine extrema of the chemical potential for this temperature and volume.
				tol = 1.0e-8;
				mu0 = SolveChemicalPotential(parent.count_min + tol, -1.0e8, 1.0e8);
				mu1 = SolveChemicalPotential(parent.count_max - tol, -1.0e8, 1.0e8);
			}

			private double LogPartitionFunction(double mu, out double[] f)
			{
				double[] e = new double[n];
				double e0 = double.NegativeInfinity;
				for (int i = 0; i < n; i++)
				{
					e[i] = beta * mu * parent.counts[i] - beta * energies[i];
					e0 = Math.Max(e0, e[i]);
				}

				// Compute partition function with prefactor introduced for numerical stability.
				double z = 0.0;
				f = new double[n];
				for (int i = 0; i < n; i++)
				{
					f[i] = Math.Exp(e[i] - e0) * parent.Entries[i].Degeneracy;
					z += f[i];
				}

				for (int i = 0; i < n; i++)
				{
					// Prefactor doesn't matter for the Boltzmann factor.
					f[i] /= z;
				}

				// Divide out prefactor again after taking the logarithm.
				double lz = Math.Log(z) + e0;

				return lz;
			}

			private void LogPartitionFunction(double mu, out double lz, out double lz_v, out double lz_mu, out double[] f, out double[] f_v, out double[] f_mu)
			{
				double[] e = new double[n];
				double[] e_v = new double[n];
				double[] e_mu = new double[n];
				double e0 = double.NegativeInfinity;
				double e0_v = double.NaN;
				double e0_mu = double.NaN;
				for (int i = 0; i < n; i++)
				{
					e[i] = beta * mu * parent.counts[i] - beta * energies[i];
					e_v[i] = beta * pressures[i];
					e_mu[i] = beta * parent.counts[i];

					if (e[i] > e0)
					{
						// Derivatives of the step function.
						e0 = e[i];
						e0_v = e_v[i];
						e0_mu = e_mu[i];
					}
				}

				// Compute partition function with prefactor introduced for numerical stability.
				double z = 0.0;
				double z_v = 0.0;
				double z_mu = 0.0;
				f = new double[n];
				f_v = new double[n];
				f_mu = new double[n];
				for (int i = 0; i < n; i++)
				{
					f[i] = Math.Exp(e[i] - e0) * parent.Entries[i].Degeneracy;
					f_v[i] = f[i] * (e_v[i] - e0_v);
					f_mu[i] = f[i] * (e_mu[i] - e0_mu);
					z += f[i];
					z_v += f_v[i];
					z_mu += f_mu[i];
				}

				for (int i = 0; i < n; i++)
				{
					// Prefactor doesn't matter for the Boltzmann factor.
					f_v[i] = f_v[i] / z - f[i] / (z * z) * z_v;
					f_mu[i] = f_mu[i] / z - f[i] / (z * z) * z_mu;
					f[i] /= z;
				}

				// Divide out prefactor again after taking the logarithm.
				lz = Math.Log(z) + e0;
				lz_v = 1.0 / z * z_v + e0_v;
				lz_mu = 1.0 / z * z_mu + e0_mu;
			}
			
			public double GrandPotential(double mu)
			{
				double phi, phi_v, phi_mu;
				GrandPotential(mu, out phi, out phi_v, out phi_mu);

				return phi;
			}

			public void GrandPotential(double mu, out double phi, out double phi_v, out double phi_mu)
			{
				double lz, lz_v, lz_mu;
				double[] f, f_v, f_mu;
				LogPartitionFunction(mu, out lz, out lz_v, out lz_mu, out f, out f_v, out f_mu);

				phi = -lz / beta;
				phi_v = -lz_v / beta;
				phi_mu = -lz_mu / beta;
			}

			public double Count(double mu)
			{
				//double x, x_v, x_mu;
				//Count(mu, out x, out x_v, out x_mu);

				double[] f;
				LogPartitionFunction(mu, out f);

				double x = 0.0;
				for (int i = 0; i < n; i++)
				{
					// Prefactor doesn't matter for the Boltzmann factor.
					x += f[i] * parent.counts[i];
				}

				return x;
			}	

			public void Count(double mu, out double x, out double x_v, out double x_mu)
			{
				double lz, lz_v, lz_mu;
				double[] f, f_v, f_mu;
				LogPartitionFunction(mu, out lz, out lz_v, out lz_mu, out f, out f_v, out f_mu);

				x = 0.0;
				x_v = 0.0;
				x_mu = 0.0;
				for (int i = 0; i < n; i++)
				{
					// Prefactor doesn't matter for the Boltzmann factor.
					x += f[i] * parent.counts[i];
					x_v += f_v[i] * parent.counts[i];
					x_mu += f_mu[i] * parent.counts[i];
				}
			}

			private double SolveChemicalPotential(double x, double mu0, double mu1)
			{
				return RootFinder.Bisection(mu => Count(mu) - x, mu0, mu1);
			}

			public double ChemicalPotential(double x)
			{
				//double mu, mu_v;
				//ChemicalPotential(x, out mu, out mu_v);

				if (x < parent.count_min || x > parent.count_max)
				{
					// This can't be reached.
					return double.NaN;
				}

				if (x < parent.count_min + tol)
				{
					return mu0;
				}

				if (x > parent.count_max - tol)
				{
					return mu1;
				}

				return SolveChemicalPotential(x, mu0, mu1);
			}

			public void ChemicalPotential(double x, out double mu, out double mu_v)
			{
				mu = ChemicalPotential(x);

				// Ignore the updated particle number.
				double x_v, x_mu;
				Count(mu, out x, out x_v, out x_mu);

				// Compute derivative implicitly from the inverse of the Jacobian of the change-of-variable mapping $(V,\mu)\mapsto (V,N(V,\mu)$.
				mu_v = -x_v / x_mu;
			}

			public double HelmholtzPotential(double x)
			{
				double a, a_v;
				HelmholtzPotential(x, out a, out a_v);

				return a;
			}

			public void HelmholtzPotential(double x, out double a, out double a_v)
			{
				// Legendre transform of the grand-potential.
				double mu, mu_v;
				ChemicalPotential(x, out mu, out mu_v);

				double phi, phi_v, phi_mu;
				GrandPotential(mu, out phi, out phi_v, out phi_mu);

				a = phi + mu * x;
				a_v = phi_v + phi_mu * mu_v + mu_v * x;
			}

			public double Pressure(double x)
			{
				double a, a_v;
				HelmholtzPotential(x, out a, out a_v);

				// Pressure is given as this derivative.
				double p = -a_v;

				return p;
			}
		}
	}
}
